import casadi as ca
import numpy as np

from .model import Model


class Mecanum(Model):

    def __init__(self):
        super().__init__()

        self._num_state = 3
        self._num_control = 3

        self.vx_max = 1.0
        self.vy_max = 0.5
        self.omega_max = np.pi / 5.0

        self._limit = np.array([
            [-self.vx_max, self.vx_max],
            [-self.vy_max, self.vy_max],
            [-self.omega_max, self.omega_max]
        ])

    @staticmethod
    def numpy_derivatives(x: np.ndarray, u: np.ndarray):
        return np.array([
            u[0] * np.cos(x[2]) - u[1] * np.sin(x[2]),
            u[0] * np.sin(x[2]) + u[1] * np.cos(x[2]),
            u[2]
        ])

    @staticmethod
    def casadi_derivatives(x: np.ndarray, u: np.ndarray):
        return ca.vertcat(u[0] * ca.cos(x[2]) - u[1] * ca.sin(x[2]), u[0] * ca.sin(x[2]) + u[1] * ca.cos(x[2]), u[2])
